/*
 *
 *  Licensed to the Apache Software Foundation (ASF) under one or more
 *  contributor license agreements.  See the NOTICE file distributed with
 *  this work for additional information regarding copyright ownership.
 *  The ASF licenses this file to You under the Apache License, Version 2.0
 *  (the "License"); you may not use this file except in compliance with
 *  the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 */

package t20240516_lineage;

import lombok.extern.slf4j.Slf4j;
import org.apache.calcite.plan.RelOptTable;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.metadata.JaninoRelMetadataProvider;
import org.apache.calcite.rel.metadata.RelColumnOrigin;
import org.apache.calcite.rel.metadata.RelMetadataQuery;
import org.apache.calcite.rel.metadata.RelMetadataQueryBase;
import org.apache.commons.collections.CollectionUtils;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.calcite.shaded.com.google.common.collect.Table;
import org.apache.flink.table.api.TableException;
import org.apache.flink.table.api.ValidationException;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.api.bridge.java.internal.StreamTableEnvironmentImpl;
import org.apache.flink.table.catalog.ObjectIdentifier;
import org.apache.flink.table.catalog.ResolvedCatalogTable;
import org.apache.flink.table.operations.CatalogSinkModifyOperation;
import org.apache.flink.table.operations.Operation;
import org.apache.flink.table.operations.ddl.CreateTableOperation;
import org.apache.flink.table.planner.operations.PlannerQueryOperation;
import org.apache.flink.table.planner.plan.metadata.FlinkDefaultRelMetadataProvider;
import org.apache.flink.table.planner.plan.schema.TableSourceTable;

import java.util.*;


@Slf4j
public class LineageContext {


    public static final String SQL_EMPTY_STR = "[\\s\\t\\n\\r]";

    private final StreamTableEnvironmentImpl tableEnv;


    private final Map<ObjectIdentifier, Map<String, String>> tableOptionMap;

    public LineageContext(StreamTableEnvironment tableEnv) {
        this.tableEnv = (StreamTableEnvironmentImpl)tableEnv;
        this.tableOptionMap = new HashMap<>();
    }


    public Tuple2<List<LineageRel>, LineageResult> analyzeLineage(String[] sqls) {

        List<LineageRel> lineageRelList = new ArrayList<>();
        /*
         * Since TableEnvironment is not thread-safe, add this sentence to solve it. Otherwise, NullPointerException
         * will appear when org.apache.calcite.rel.metadata.RelMetadataQuery.<init>
         * http://apache-flink.370.s1.nabble.com/flink1-11-0-sqlQuery-NullPointException-td5466.html
         */
        RelMetadataQueryBase.THREAD_PROVIDERS
                .set(JaninoRelMetadataProvider.of(FlinkDefaultRelMetadataProvider.INSTANCE()));

        for (String sql : sqls) {
            System.out.println("analyze sql:" + sql);
            SqlType operationType = getOperationType(sql);
            if (operationType.equals(SqlType.INSERT)) {
                // 1. Generate original relNode tree
                Tuple2<String, RelNode> parsed = parseStatement(sql);
                String sinkTable = parsed.getField(0);
                RelNode oriRelNode = parsed.getField(1);

                // 2. Build lineage based from RelMetadataQuery
                List<LineageRel> lineageRels = buildFiledLineageResult(sinkTable, oriRelNode);
                lineageRelList.addAll(lineageRels);

            } else if (!operationType.equals(SqlType.SELECT) && !operationType.equals(SqlType.PRINT)) {
                tableEnv.executeSql(sql);
                if (operationType.equals(SqlType.CREATE)) {
                    Operation operation = tableEnv.getParser().parse(sql).get(0);
                    if (operation instanceof CreateTableOperation) {
                        CreateTableOperation createTableOperation = (CreateTableOperation) operation;
                        Map<String, String> options = createTableOperation.getCatalogTable().getOptions();
                        ObjectIdentifier tableIdentifier = createTableOperation.getTableIdentifier();
                        tableOptionMap.put(tableIdentifier, options);
                    }
                }

            } else {
                System.err.println("not support:" + sql);
            }
        }
        return Tuple2.of(lineageRelList, buildLineage(lineageRelList));
    }

    public Tuple2<List<LineageRel>, LineageResult> analyzeLineage(List<SingleSql> sqls) {

        List<LineageRel> lineageRelList = new ArrayList<>();
        /*
         * Since TableEnvironment is not thread-safe, add this sentence to solve it. Otherwise, NullPointerException
         * will appear when org.apache.calcite.rel.metadata.RelMetadataQuery.<init>
         * http://apache-flink.370.s1.nabble.com/flink1-11-0-sqlQuery-NullPointException-td5466.html
         */
        RelMetadataQueryBase.THREAD_PROVIDERS
                .set(JaninoRelMetadataProvider.of(FlinkDefaultRelMetadataProvider.INSTANCE()));

        for (SingleSql ssql : sqls) {
            String sql = ssql.getSqlSource();
            System.out.println("analyze sql:" + sql);
            SqlType operationType = getOperationType(sql);
            if (operationType.equals(SqlType.INSERT)) {
                // 1. Generate original relNode tree
                Tuple2<String, RelNode> parsed = parseStatement(sql);
                String sinkTable = parsed.getField(0);
                RelNode oriRelNode = parsed.getField(1);

                // 2. Build lineage based from RelMetadataQuery
                List<LineageRel> lineageRels = buildFiledLineageResult(sinkTable, oriRelNode);
                lineageRelList.addAll(lineageRels);

            } else if (!operationType.equals(SqlType.SELECT) && !operationType.equals(SqlType.PRINT)) {
                tableEnv.executeSql(sql);
                if (operationType.equals(SqlType.CREATE)) {
                    Operation operation = tableEnv.getParser().parse(sql).get(0);
                    if (operation instanceof CreateTableOperation) {
                        CreateTableOperation createTableOperation = (CreateTableOperation) operation;
                        Map<String, String> options = createTableOperation.getCatalogTable().getOptions();
                        ObjectIdentifier tableIdentifier = createTableOperation.getTableIdentifier();
                        tableOptionMap.put(tableIdentifier, options);
                    }
                }

            } else {
                System.err.println("not support:" + sql);
            }
        }
        return Tuple2.of(lineageRelList, buildLineage(lineageRelList));
    }

    private Tuple2<String, RelNode> parseStatement(String sql) {
        List<Operation> operations = tableEnv.getParser().parse(sql);

        if (operations.size() != 1) {
            throw new TableException("Unsupported SQL query! only accepts a single SQL statement.");
        }
        Operation operation = operations.get(0);
        System.out.println("operation:" + operation);
        if (operation instanceof CatalogSinkModifyOperation) {
            CatalogSinkModifyOperation sinkOperation = (CatalogSinkModifyOperation) operation;

            PlannerQueryOperation queryOperation = (PlannerQueryOperation) sinkOperation.getChild();
            RelNode relNode = queryOperation.getCalciteTree();
            return new Tuple2<>(sinkOperation.getTableIdentifier().asSummaryString(), relNode);
        } else {
            throw new TableException("Only insert is supported now.");
        }
    }

    /** Check the size of query and sink fields match */
    private void validateSchema(String sinkTable, RelNode relNode, List<String> sinkFieldList) {
        List<String> queryFieldList = relNode.getRowType().getFieldNames();
        if (queryFieldList.size() != sinkFieldList.size()) {
            throw new ValidationException(String.format(
                    "Column types of query result and sink for %s do not match.\n"
                            + "Query schema: %s\n"
                            + "Sink schema:  %s",
                    sinkTable, queryFieldList, sinkFieldList));
        }
    }

    private List<LineageRel> buildFiledLineageResult(String sinkTable, RelNode optRelNode) {
        // target columns
        List<String> targetColumnList =
                tableEnv.from(sinkTable).getResolvedSchema().getColumnNames();

        // check the size of query and sink fields match
        validateSchema(sinkTable, optRelNode, targetColumnList);

        RelMetadataQuery metadataQuery = optRelNode.getCluster().getMetadataQuery();
        List<LineageRel> resultList = new ArrayList<>();



//        System.out.println("----------------------------------------------------------------------");
//        Table<RelNode, List, Object> map = metadataQuery.map;
//        for (Table.Cell<RelNode, List, Object> relNodeListObjectCell : map.cellSet()) {
//            RelNode rowKey = relNodeListObjectCell.getRowKey();
//            String tbName = String.join(",", rowKey.getTable().getQualifiedName());
//
//            System.out.println("tbName:" + tbName);
//            List columnKey = relNodeListObjectCell.getColumnKey();
//            for (Object o : columnKey) {
//                System.out.println("  " + o);
//            }
//            Object value = relNodeListObjectCell.getValue();
//            System.out.println("    value:" + value);
//        }
//        System.out.println("----------------------------------------------------------------------");

        for (int index = 0; index < targetColumnList.size(); index++) {
            String targetColumn = targetColumnList.get(index);


            Set<RelColumnOrigin> relColumnOriginSet = metadataQuery.getColumnOrigins(optRelNode, index);

            System.out.println(String.format("targetColumn:%s from sourceCount:%d", targetColumn, relColumnOriginSet.size()));
            if (CollectionUtils.isNotEmpty(relColumnOriginSet)) {
                for (RelColumnOrigin relColumnOrigin : relColumnOriginSet) {
                    // table
                    RelOptTable table = relColumnOrigin.getOriginTable();
                    String sourceTable = String.join(".", table.getQualifiedName());

                    // filed
                    int ordinal = relColumnOrigin.getOriginColumnOrdinal();
                    ResolvedCatalogTable resolvedCatalogTable = ((TableSourceTable) table).catalogTable();

                    List<String> columnNames = resolvedCatalogTable.getResolvedSchema().getColumnNames();
                    String sourceColumn = columnNames.get(ordinal);

                    // add record
                    resultList.add(LineageRel.build(
                            sourceTable, sourceColumn,
                            sinkTable, targetColumn,
                            relColumnOrigin.getTransform()));
                }
            }
        }
        return resultList;
    }


    public LineageResult buildLineage(List<LineageRel> lineageRelList) {
        List<LineageRelation> relations = new ArrayList<>();
        Map<String, LineageTable> tableMap = new HashMap<>();
        int tableIndex = 1;
        int relIndex = 1;
        for (LineageRel lineageRel : lineageRelList) {
            String sourceTablePath = lineageRel.getSourceTablePath();
            String sourceTableId = null;
            String targetTableId = null;
            if (tableMap.containsKey(sourceTablePath)) {
                LineageTable lineageTable = tableMap.get(sourceTablePath);
                LineageColumn lineageColumn =
                        LineageColumn.build(lineageRel.getSourceColumn(), lineageRel.getSourceColumn());
                if (!lineageTable.getColumns().contains(lineageColumn)) {
                    lineageTable.getColumns().add(lineageColumn);
                }
                sourceTableId = lineageTable.getId();
            } else {
                tableIndex++;
                LineageTable lineageTable = LineageTable.build(String.valueOf(tableIndex), sourceTablePath,
                        tableOptionMap.get(lineageRel.getSourceTableIdentifier()).get("connector"));
                lineageTable
                        .getColumns()
                        .add(LineageColumn.build(lineageRel.getSourceColumn(), lineageRel.getSourceColumn()));
                tableMap.put(sourceTablePath, lineageTable);
                sourceTableId = lineageTable.getId();
            }
            String targetTablePath = lineageRel.getTargetTablePath();
            if (tableMap.containsKey(targetTablePath)) {
                LineageTable lineageTable = tableMap.get(targetTablePath);
                LineageColumn lineageColumn = LineageColumn.build(lineageRel.getTargetColumn(), lineageRel.getTargetColumn());
                if (!lineageTable.getColumns().contains(lineageColumn)) {
                    lineageTable.getColumns().add(lineageColumn);
                }
                targetTableId = lineageTable.getId();
            } else {
                tableIndex++;
                LineageTable lineageTable = LineageTable.build(String.valueOf(tableIndex), targetTablePath,
                        tableOptionMap.get(lineageRel.getTargetTableIdentifier()).get("connector"));
                lineageTable
                        .getColumns()
                        .add(LineageColumn.build(lineageRel.getTargetColumn(), lineageRel.getTargetColumn()));
                tableMap.put(targetTablePath, lineageTable);
                targetTableId = lineageTable.getId();
            }
            LineageRelation lineageRelation = LineageRelation.build(
                    sourceTableId, targetTableId, lineageRel.getSourceColumn(), lineageRel.getTargetColumn());
            if (!relations.contains(lineageRelation)) {
                relIndex++;
                lineageRelation.setId(relIndex + "");
                relations.add(lineageRelation);
            }
        }
        List<LineageTable> tables = new ArrayList<LineageTable>(tableMap.values());
        return LineageResult.build(tables, relations);
    }


    public static SqlType getOperationType(String sql) {
        String sqlTrim = sql.replaceAll(SQL_EMPTY_STR, " ").trim().toUpperCase();
        return Arrays.stream(SqlType.values())
                .filter(sqlType -> sqlType.match(sqlTrim))
                .findFirst()
                .orElse(SqlType.UNKNOWN);
    }

}
